# from modelscope import snapshot_download
# model_dir = snapshot_download("AI-ModelScope/bge-large-zh-v1.5", revision='master')
# print(model_dir)
# from langchain_community.embeddings import HuggingFaceBgeEmbeddings
# model_name = model_dir
# model_kwargs = {'device': 'cuda'}
# encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
# model = HuggingFaceBgeEmbeddings(
#     model_name=model_name,
#     model_kwargs=model_kwargs,
#     encode_kwargs=encode_kwargs,
#     query_instruction="为这个句子生成表示以用于检索相关文章："
# )
# model.query_instruction = "为这个句子生成表示以用于检索相关文章："
#
# text = "这是一个测试句子。"
# embeddings = model.embed_query(text)
# print(embeddings)


import os
from langchain_community.chat_models import ChatZhipuAI

from langchain_community.tools import DuckDuckGoSearchRun
from langchain.tools import tool
from langchain_core.prompts import PromptTemplate
from langchain_core.prompts import FewShotPromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline

def example_select_demo():
    os.environ["ZHIPUAI_API_KEY"] = "544d7293332c56b92739130d933f115b.x4GONk9O2pkdJFDh"

    # 初始化ChatZhipuAI模型
    zhipuai = ChatZhipuAI(
        model="glm-4",
        temperature=0.5,
        ###
        ###
    )


    # 初始化embding模型
    # 创建嵌入生成器和向量数据库
    embedder = HuggingFaceEmbeddings(
        model_name="/data/ai/langchain/sentence_embedding_model/AI-ModelScope/all-MiniLM-L6-v2"
    )
    example_prompt = PromptTemplate.from_template("Question: {question}\n{answer}")

    examples = [
        {
            "question": "Craigslist的创始人出生于何时？",
            "answer": """
    这里是否需要后续问题：是。
    后续问题：谁是Craigslist的创始人？
    中间答案：Craigslist的创始人是克雷格·纽马克。
    后续问题：克雷格·纽马克是什么时候出生的？
    中间答案：克雷格·纽马克出生于1952年12月6日。
    因此最终答案是：1952年12月6日
    """,
        },
        {
            "question": "乔治·华盛顿的外祖父是谁？",
            "answer": """
    这里是否需要后续问题：是。
    后续问题：乔治·华盛顿的母亲是谁？
    中间答案：乔治·华盛顿的母亲是玛丽·鲍尔·华盛顿。
    后续问题：玛丽·鲍尔·华盛顿的父亲是谁？
    中间答案：玛丽·鲍尔·华盛顿的父亲是约瑟夫·鲍尔。
    因此最终答案是：约瑟夫·鲍尔
    """,
        },
    ]

    # 打印示例
    # print(example_prompt.invoke(examples[0]).to_string())

    prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
        suffix="Question: {input}",
        input_variables=["input"],
    )


    # prompt样例
    # print(
    #     prompt.invoke({"input": "腾讯的创始人出生于何时"}).to_string()
    # )

    from langchain_chroma import Chroma
    from langchain_core.example_selectors import SemanticSimilarityExampleSelector
    from langchain_openai import OpenAIEmbeddings

    example_selector = SemanticSimilarityExampleSelector.from_examples(
        # This is the list of examples available to select from.
        examples,
        # This is the embedding class used to produce embeddings which are used to measure semantic similarity.
        embedder,
        # This is the VectorStore class that is used to store the embeddings and do a similarity search over.
        Chroma,
        # This is the number of examples to produce.
        k=1,
    )

    # Select the most similar example to the input.
    question = "腾讯的创始人生于什么时间"
    selected_examples = example_selector.select_examples({"question": question})
    print(f"Examples most similar to the input: {question}")
    for example in selected_examples:
        print("\n")
        for k, v in example.items():
            print(f"{k}: {v}")




if __name__ == '__main__':
    # 示例选择demo
    example_select_demo()